import requests
from transformers import AutoTokenizer

class PostSender:
    def __init__(self, post_url, model_dir):
        self.post_url = post_url
        self.model_dir = model_dir

    @staticmethod
    def scale_prompt_length(prompt, model_dir):
        # 定义一个区间表
        intervals = [128, 256, 512, 1024, 2048]
        model_dir = model_dir.lower()
        if "llama" in model_dir:
            # length为文本长度
            length = len(prompt)
            # 计算prompt tokenizer后的长度
            tokenizer = AutoTokenizer.from_pretrained("/home/user1/LLaMA-Factory/Llama-2-7b-hf/", token="HF_TOKEN")
            tokens = tokenizer.encode(prompt)

            # 如果tokens的长度大于2200，直接不要
            # 如果tokens的长度大于2048，那么直接将tokens截断到2048，并且返回解码后的字符串
            if 2048 < len(tokens) < 2200:
                return tokenizer.decode(tokens[:2047]), 2048
            else:
                max_tokens = next((x for x in intervals if x >= length), 2048)
                return prompt, max_tokens

        else:
            length = len(prompt)
            scaled_length = length / 1.5
            # 定义一个区间表
            intervals = [128, 256, 512, 1024, 2048]
            # 找到第一个大于scaled_length的区间右端点
            max_tokens = next((x for x in intervals if x >= scaled_length), 2048)
            return prompt,max_tokens

    def send(self, prompt) -> str:
        # 对prompt做缩放处理，首先将prompt除以1.5, 查看结果落在哪个区间，(0,128],(128, 256], (256, 512], (512, 1024], (1024, 2048], 取所在区间的右端点
        processed_prompt, max_tokens = PostSender.scale_prompt_length(prompt, self.model_dir)

        # 注意utf-8编码
        body = {
            "model": self.model_dir,
            "prompt": processed_prompt,
            "max_tokens": max_tokens,
            "temperature": 0.4,
            "top_p": 0.9,
            "repetition_penalty": 1.15
        }
        headers = {
            'Content-Type': 'application/json'
        }

        try:
            response = requests.post(self.post_url, json=body, headers=headers)
            response_json = response.json()
            return response_json.get("choices")[0].get("text")
        except:
            # 打印错误信息
            print("Error: ", response.text)

if __name__ == '__main__':  
    post_url = "http://localhost:8000/v1/completions"
    post_sender = PostSender(post_url, "/home/user1/models/formal/Llama_tuned/")
    prompt = "你将完成生成中国司法裁判文书中刑事一审判决书的任务，该类文书分为以下几段：证据段、分析段、引用法条段、判决段。接下来将提供证据段，请你仅生成后续的分析段。证据段如下:南京市秦淮区人民检察院指控，2016年2月21日20时许，被告人巴某与麦某（另案处理）在本市秦淮区许家巷44号秦淮医院门口处，趁被害人何某乙不备，窃得何某乙上衣右侧口袋内的白色苹果牌iPhone5S型手机一部，价值人民币1100元，后被公安机关当场抓获。被告人巴某归案后如实供述了上述犯罪事实。案发后，被盗的白色苹果牌iPhone5S型手机一部被公安机关依法追缴并发还被害人何某乙。上述事实，被告人巴某在开庭审理过程中亦无异议，且有被害人何某乙的陈述、证人麦某、许某、严某、何某甲等人的证言、扣押及发还物品清单、刑事摄影照片、搜查笔录、价格鉴证结论书、前科材料、情况说明、抓获经过、户籍资料等证据证实，足以认定。"
    print(post_sender.send(prompt))

